import pandas as pd
import numpy as np

print("⏳ 正在读取文件...")
# 1. 读取原始数据
df = pd.read_parquet('exercise_contest_train-00000-of-00001.parquet')

# 打印一下真正的列名，确认一下 (调试用)
print("📋 文件包含的列名：", df.columns.tolist())

print("🧹 开始清洗数据...")

# 2. 只需要重命名一下列，把它改成 App 能识别的名字
# 我们创建一个新的干净 DataFrame
final_df = pd.DataFrame()

# 复制案情
final_df['fact'] = df['fact']

# 处理罪名 (accusation 在数据里通常是一个列表 ['故意伤害']，我们要取出里面的文字)
# 这一行代码会自动取列表的第一个元素
def get_first_charge(x):
    try:
        # 如果是列表或数组，取第一个
        if isinstance(x, (list, np.ndarray)) and len(x) > 0:
            return x[0]
        return str(x)
    except:
        return "未知"

final_df['real_charge'] = df['accusation'].apply(get_first_charge)

# 处理刑期 (直接用 imprisonment 列)
final_df['real_term'] = df['imprisonment'].fillna(0).astype(int)

# 3. 过滤一下数据
# 只保留有刑期的案子 (刑期 > 0) 并且只取前 20000 条做 Demo
final_df = final_df[final_df['real_term'] > 0].head(20000)

print(f"🎉 清洗完成！有效数据剩余：{len(final_df)} 条")
print("数据样例：")
print(final_df.head(3))

# 4. 保存为 Pickle
final_df.to_pickle('law_data_clean.pkl')
print("💾 已成功保存为 'law_data_clean.pkl'")